"""系统启动时的任务状态检查服务
负责在系统启动时检查和修正任务状态
"""
import logging
import traceback
from datetime import datetime, timedelta
from typing import List, Dict, Any

from services.shared.db_pool import get_db_session, close_db_session
from services.shared.models.task import TaskExecution, Task
from services.shared.models.enums import TaskStatus
from services.shared.models.system import SystemLog
from services.shared.celery_app import celery_app
from services.shared.timezone_utils import now as timezone_now, naive_to_aware
from services.shared.log_utils import log_system_event, log_error, log_info
from services.shared.logger import get_logger

# 获取日志记录器
logger = get_logger('shared.startup_task_checker')


class StartupTaskChecker:
    """系统启动时的任务状态检查器"""
    
    def __init__(self):
        self.cleanup_records = []
        
    def check_and_fix_task_status(self) -> Dict[str, Any]:
        """检查并修正任务状态
        
        Returns:
            检查结果字典
        """
        logger.info("开始系统启动时的任务状态检查")
        
        start_time = timezone_now()
        result = {
            'success': True,
            'start_time': start_time.isoformat(),
            'pending_checked': 0,
            'pending_fixed': 0,
            'running_checked': 0,
            'running_fixed': 0,
            'errors': []
        }
        
        try:
            # 检查等待中的任务
            pending_result = self._check_pending_tasks()
            result['pending_checked'] = pending_result['checked_count']
            result['pending_fixed'] = pending_result['fixed_count']
            
            # 检查运行中的任务
            running_result = self._check_running_tasks()
            result['running_checked'] = running_result['checked_count']
            result['running_fixed'] = running_result['fixed_count']
            
            # 保存清理记录
            if self.cleanup_records:
                self._save_cleanup_records()
            
            result['end_time'] = timezone_now().isoformat()
            result['duration'] = (timezone_now() - start_time).total_seconds()
            
            # 记录系统事件
            log_system_event(
                level='INFO',
                module='startup_checker',
                message='系统启动任务状态检查完成',
                details={
                    'checked_pending': result['pending_checked'],
                    'fixed_pending': result['pending_fixed'],
                    'checked_running': result['running_checked'],
                    'fixed_running': result['running_fixed'],
                    'duration': result['duration']
                }
            )
            
            logger.info(f"任务状态检查完成: 等待中任务检查{result['pending_checked']}个，修正{result['pending_fixed']}个；运行中任务检查{result['running_checked']}个，修正{result['running_fixed']}个")
            
        except Exception as e:
            error_msg = f"任务状态检查失败: {str(e)}"
            logger.error(error_msg)
            logger.debug(f"异常详情: {traceback.format_exc()}")
            
            result['success'] = False
            result['errors'].append(error_msg)
            
            # 记录错误事件
            log_error(
                "startup_task_check_failed",
                {
                    "error_type": type(e).__name__,
                    "error_message": str(e),
                    "stack_trace": traceback.format_exc()
                }
            )
        
        return result
    
    def _check_pending_tasks(self) -> Dict[str, int]:
        """检查等待中的任务
        
        Returns:
            检查结果字典
        """
        logger.debug("开始检查等待中的任务")
        
        session = get_db_session()
        checked_count = 0
        fixed_count = 0
        
        try:
            # 查询所有等待中的任务执行记录
            pending_executions = session.query(TaskExecution).filter(
                TaskExecution.status == TaskStatus.PENDING
            ).all()
            
            checked_count = len(pending_executions)
            logger.debug(f"找到 {checked_count} 个等待中的任务")
            
            if not pending_executions:
                return {'checked_count': 0, 'fixed_count': 0}
            
            # 获取Celery队列中的任务
            celery_pending_ids = self._get_celery_pending_tasks()
            logger.debug(f"Celery队列中有 {len(celery_pending_ids)} 个等待任务")
            
            # 检查每个等待中的任务
            for execution in pending_executions:
                if self._should_fix_pending_task(execution, celery_pending_ids):
                    self._fix_pending_task(execution, session)
                    fixed_count += 1
            
            session.commit()
            logger.info(f"等待中任务检查完成: 检查 {checked_count} 个，修正 {fixed_count} 个")
            
        except Exception as e:
            logger.error(f"检查等待中任务失败: {str(e)}")
            session.rollback()
            raise
        finally:
            session.close()
        
        return {'checked_count': checked_count, 'fixed_count': fixed_count}
    
    def _check_running_tasks(self) -> Dict[str, int]:
        """检查运行中的任务
        
        Returns:
            检查结果字典
        """
        logger.debug("开始检查运行中的任务")
        
        session = get_db_session()
        checked_count = 0
        fixed_count = 0
        
        try:
            # 查询所有运行中的任务执行记录
            running_executions = session.query(TaskExecution).filter(
                TaskExecution.status == TaskStatus.RUNNING
            ).all()
            
            checked_count = len(running_executions)
            logger.debug(f"找到 {checked_count} 个运行中的任务")
            
            if not running_executions:
                return {'checked_count': 0, 'fixed_count': 0}
            
            # 获取Celery中活跃的任务
            celery_active_ids = self._get_celery_active_tasks()
            logger.debug(f"Celery中有 {len(celery_active_ids)} 个活跃任务")
            
            # 检查每个运行中的任务
            for execution in running_executions:
                if self._should_fix_running_task(execution, celery_active_ids):
                    # 尝试撤销Celery任务
                    if hasattr(execution, 'celery_task_id') and execution.celery_task_id:
                        try:
                            celery_app.control.revoke(execution.celery_task_id, terminate=True)
                            logger.debug(f"已撤销Celery任务: {execution.celery_task_id}")
                        except Exception as revoke_error:
                            logger.warning(f"撤销Celery任务失败: {revoke_error}")
                    
                    self._fix_running_task(execution, session)
                    fixed_count += 1
            
            session.commit()
            logger.info(f"运行中任务检查完成: 检查 {checked_count} 个，修正 {fixed_count} 个")
            
        except Exception as e:
            logger.error(f"检查运行中任务失败: {str(e)}")
            session.rollback()
            raise
        finally:
            session.close()
        
        return {'checked_count': checked_count, 'fixed_count': fixed_count}
    
    def _get_celery_pending_tasks(self) -> set:
        """获取Celery队列中等待的任务ID
        
        Returns:
            等待任务的execution_id集合
        """
        celery_pending_ids = set()
        
        try:
            inspect = celery_app.control.inspect()
            
            # 获取预留任务（等待执行的任务）
            reserved_tasks = inspect.reserved()
            if reserved_tasks:
                for worker, tasks in reserved_tasks.items():
                    if not tasks:
                        continue
                    
                    for task in tasks:
                        if task.get('name') == 'services.executor.tasks.execute_task':
                            execution_id = self._extract_execution_id_from_task(task)
                            if execution_id:
                                celery_pending_ids.add(execution_id)
            
            # 获取调度任务（即将执行的任务）
            scheduled_tasks = inspect.scheduled()
            if scheduled_tasks:
                for worker, tasks in scheduled_tasks.items():
                    if not tasks:
                        continue
                    
                    for task in tasks:
                        if task.get('name') == 'services.executor.tasks.execute_task':
                            execution_id = self._extract_execution_id_from_task(task)
                            if execution_id:
                                celery_pending_ids.add(execution_id)
        
        except Exception as e:
            logger.warning(f"获取Celery等待任务失败: {str(e)}")
            logger.debug(f"异常详情: {traceback.format_exc()}")
        
        return celery_pending_ids
    
    def _get_celery_active_tasks(self) -> set:
        """获取Celery中活跃的任务ID
        
        Returns:
            活跃任务的execution_id和celery_task_id集合
        """
        celery_active_ids = set()
        
        try:
            inspect = celery_app.control.inspect()
            
            # 获取活跃任务（正在执行的任务）
            active_tasks = inspect.active()
            if active_tasks:
                for worker, tasks in active_tasks.items():
                    if not tasks:
                        continue
                    
                    for task in tasks:
                        # 添加Celery任务ID
                        celery_active_ids.add(task.get('id'))
                        
                        # 如果是执行任务，提取execution_id
                        if task.get('name') == 'services.executor.tasks.execute_task':
                            execution_id = self._extract_execution_id_from_task(task)
                            if execution_id:
                                celery_active_ids.add(execution_id)
            
            # 获取预定任务
            scheduled_tasks = inspect.scheduled()
            if scheduled_tasks:
                for worker, tasks in scheduled_tasks.items():
                    if not tasks:
                        continue
                    
                    for task in tasks:
                        # 添加Celery任务ID
                        task_id = task.get('request', {}).get('id') or task.get('id')
                        if task_id:
                            celery_active_ids.add(task_id)
                        
                        # 如果是执行任务，提取execution_id
                        if task.get('name') == 'services.executor.tasks.execute_task':
                            execution_id = self._extract_execution_id_from_task(task)
                            if execution_id:
                                celery_active_ids.add(execution_id)
            
            # 获取保留任务
            reserved_tasks = inspect.reserved()
            if reserved_tasks:
                for worker, tasks in reserved_tasks.items():
                    if not tasks:
                        continue
                    
                    for task in tasks:
                        # 添加Celery任务ID
                        celery_active_ids.add(task.get('id'))
                        
                        # 如果是执行任务，提取execution_id
                        if task.get('name') == 'services.executor.tasks.execute_task':
                            execution_id = self._extract_execution_id_from_task(task)
                            if execution_id:
                                celery_active_ids.add(execution_id)
        
        except Exception as e:
            logger.warning(f"获取Celery活跃任务失败: {str(e)}")
            logger.debug(f"异常详情: {traceback.format_exc()}")
        
        return celery_active_ids
    
    def _extract_execution_id_from_task(self, task: Dict[str, Any]) -> str:
        """从Celery任务信息中提取execution_id
        
        Args:
            task: Celery任务信息字典
            
        Returns:
            execution_id或None
        """
        try:
            args = task.get('args', [])
            kwargs = task.get('kwargs', {})
            
            # 方式1: 从kwargs中的task_data获取
            if 'task_data' in kwargs:
                task_data = kwargs['task_data']
                if isinstance(task_data, dict):
                    return task_data.get('execution_id')
            
            # 方式2: 从kwargs直接获取
            if 'execution_id' in kwargs:
                return kwargs['execution_id']
            
            # 方式3: 从args中获取
            if args and isinstance(args[0], dict):
                return args[0].get('execution_id')
            
            # 方式4: 从args的第二个参数获取
            if len(args) >= 2:
                return args[1]
        
        except Exception as e:
            logger.debug(f"提取execution_id失败: {str(e)}")
        
        return None
    
    def _should_fix_pending_task(self, execution: TaskExecution, celery_pending_ids: set) -> bool:
        """判断是否需要修正等待中的任务
        
        Args:
            execution: 任务执行记录
            celery_pending_ids: Celery队列中等待的任务ID集合
            
        Returns:
            是否需要修正
        """
        # 如果任务在Celery队列中，不需要修正
        if execution.execution_id in celery_pending_ids:
            return False
        
        # 检查任务是否等待太久（超过10分钟可能是孤立任务）
        if execution.started_at:
            current_time = timezone_now()
            # 确保时间比较的一致性
            started_at = execution.started_at
            
            # 确保两个时间都是带时区的或都是naive的
            if started_at.tzinfo is None:
                started_at = naive_to_aware(started_at)
            elif current_time.tzinfo is None:
                current_time = naive_to_aware(current_time)
            
            wait_duration = current_time - started_at
            if wait_duration > timedelta(minutes=10):
                logger.debug(f"任务 {execution.execution_id} 等待时间过长: {wait_duration}")
                return True
        
        return False
    
    def _should_fix_running_task(self, execution: TaskExecution, celery_active_ids: set) -> bool:
        """判断是否需要修正运行中的任务
        
        Args:
            execution: 任务执行记录
            celery_active_ids: Celery中活跃的任务ID集合
            
        Returns:
            是否需要修正
        """
        # 如果任务在Celery中活跃，不需要修正
        if execution.execution_id in celery_active_ids:
            return False
        
        # 如果有Celery任务ID，检查该任务是否还存在
        if hasattr(execution, 'celery_task_id') and execution.celery_task_id:
            # 检查任务ID是否在活跃任务列表中
            if execution.celery_task_id in celery_active_ids:
                return False
        
        # 优先检查进程是否还在运行（这是最可靠的判断方式）
        if hasattr(execution, 'executor_pid') and execution.executor_pid:
            try:
                import os
                import signal
                os.kill(execution.executor_pid, 0)  # 发送信号0检查进程是否存在
                return False  # 进程存在，不需要修正
            except (OSError, ProcessLookupError):
                logger.debug(f"进程 {execution.executor_pid} 不存在，需要修正任务")
                return True  # 进程不存在，需要修正
        
        # 如果没有进程ID，再检查Celery任务状态
        if hasattr(execution, 'celery_task_id') and execution.celery_task_id:
            try:
                # 尝试获取任务状态
                from celery.result import AsyncResult
                result = AsyncResult(execution.celery_task_id, app=celery_app)
                if result.state in ['PENDING', 'STARTED', 'RETRY']:
                    # 如果Celery显示任务还在运行，但没有进程ID，这是异常情况
                    logger.debug(f"任务 {execution.execution_id} Celery状态为 {result.state} 但没有进程ID")
                    return True
            except Exception as e:
                logger.debug(f"检查Celery任务状态失败: {str(e)}")
        
        # 检查任务运行时间是否过长（超过24小时视为异常）
        if execution.started_at:
            current_time = timezone_now()
            started_at = execution.started_at
            
            # 确保两个时间都是带时区的或都是naive的
            if started_at.tzinfo is None:
                started_at = naive_to_aware(started_at)
            elif current_time.tzinfo is None:
                current_time = naive_to_aware(current_time)
            
            running_duration = current_time - started_at
            if running_duration.total_seconds() > 24 * 3600:  # 24小时
                logger.debug(f"任务 {execution.execution_id} 运行时间过长: {running_duration.total_seconds()/3600:.1f}小时")
                return True
        
        return True
    
    def _fix_pending_task(self, execution: TaskExecution, session) -> None:
        """修正等待中的孤立任务
        
        Args:
            execution: 任务执行记录
            session: 数据库会话
        """
        logger.info(f"修正等待中的孤立任务: {execution.execution_id}")
        
        # 更新任务状态为失败
        execution.status = TaskStatus.FAILED
        execution.completed_at = timezone_now()
        execution.error_message = '系统启动时检测到任务未在执行器队列中，已自动标记为失败'
        
        # 计算执行时长
        if execution.started_at:
            started_at = execution.started_at
            completed_at = execution.completed_at
            
            # 确保两个时间都是带时区的或都是naive的
            if started_at.tzinfo is None:
                started_at = naive_to_aware(started_at)
            elif completed_at.tzinfo is None:
                completed_at = naive_to_aware(completed_at)
            
            duration = (completed_at - started_at).total_seconds()
            execution.duration = max(0, duration)
        
        # 记录清理信息
        self.cleanup_records.append({
            'execution_id': execution.execution_id,
            'task_id': execution.task_id,
            'original_status': 'PENDING',
            'new_status': 'FAILED',
            'reason': '任务未在执行器队列中',
            'fixed_at': timezone_now().isoformat()
        })
        
        logger.debug(f"等待中任务 {execution.execution_id} 已修正为失败状态")
    
    def _fix_running_task(self, execution: TaskExecution, session) -> None:
        """修正运行中的孤立任务
        
        Args:
            execution: 任务执行记录
            session: 数据库会话
        """
        logger.info(f"修正运行中的孤立任务: {execution.execution_id}")
        
        # 更新任务状态为失败
        execution.status = TaskStatus.FAILED
        execution.completed_at = timezone_now()
        execution.error_message = '系统启动时检测到任务未在真实执行中，已自动标记为失败'
        
        # 计算执行时长
        if execution.started_at:
            started_at = execution.started_at
            completed_at = execution.completed_at
            
            # 确保两个时间都是带时区的或都是naive的
            if started_at.tzinfo is None:
                started_at = naive_to_aware(started_at)
            elif completed_at.tzinfo is None:
                completed_at = naive_to_aware(completed_at)
            
            duration = (completed_at - started_at).total_seconds()
            execution.duration = max(0, duration)
        
        # 记录清理信息
        self.cleanup_records.append({
            'execution_id': execution.execution_id,
            'task_id': execution.task_id,
            'original_status': 'RUNNING',
            'new_status': 'FAILED',
            'reason': '任务未在真实执行中',
            'fixed_at': timezone_now().isoformat()
        })
        
        logger.debug(f"运行中任务 {execution.execution_id} 已修正为失败状态")
    
    def _save_cleanup_records(self) -> None:
        """保存清理记录到系统日志"""
        if not self.cleanup_records:
            return
        
        try:
            session = get_db_session()
            
            # 根据修正数量确定日志级别
            fixed_count = len(self.cleanup_records)
            level = 'WARNING' if fixed_count > 0 else 'INFO'
            
            # 创建系统日志记录
            log_entry = SystemLog(
                level=level,
                module='startup_checker',
                message=f'系统启动时修正了 {fixed_count} 个孤立任务',
                details={
                    'cleanup_type': 'startup_task_status_fix',
                    'fixed_count': fixed_count,
                    'records': self.cleanup_records,
                    'timestamp': timezone_now().isoformat()
                },
                created_at=timezone_now()
            )
            
            session.add(log_entry)
            session.commit()
            
            if fixed_count > 0:
                logger.warning(f"已保存 {fixed_count} 条任务修正记录到系统日志")
            else:
                logger.info(f"已保存 {fixed_count} 条清理记录")
            
        except Exception as e:
            logger.error(f"保存清理记录失败: {str(e)}")
            logger.debug(f"异常详情: {traceback.format_exc()}")
            # 尝试使用备用日志记录方式
            try:
                log_error(
                    "startup_cleanup_record_save_failed",
                    {
                        "error_type": type(e).__name__,
                        "error_message": str(e),
                        "cleanup_records_count": len(self.cleanup_records)
                    }
                )
            except Exception:
                pass
        finally:
            close_db_session(session)


def run_startup_task_check() -> Dict[str, Any]:
    """运行系统启动时的任务状态检查
    
    Returns:
        检查结果字典
    """
    checker = StartupTaskChecker()
    return checker.check_and_fix_task_status()


if __name__ == '__main__':
    # 直接运行检查
    result = run_startup_task_check()
    from services.shared.log_utils import log_info
    import json
    log_info(
        module='startup_task_checker',
        message="任务状态检查完成",
        details=json.dumps(result, ensure_ascii=False)
    )
    logger.info(f"任务状态检查结果: {result}")